#!/usr/bin/env python
"""

   translate_by_codon.py

        accepts cDNA sequence and codon by codon translation.
        Unknown codons are translated as 'X'
 

"""

import sys


#_________________________________________________________________________________________

#   Make maps

#_________________________________________________________________________________________
stops = dict.fromkeys(['TGA', 'TAA', 'TAG'])
def four_fold_deg (codon_to_aa_map, codon12, aa):
    for third in 'A', 'C', 'G', 'T', 'X', 'N', '-', ' ':
        codon_to_aa_map[codon12 + third] = aa
def create_codon_to_aa_map():
    codon_to_aa_map = dict()


    # stops
    codon_to_aa_map['TGA'] = '*'
    codon_to_aa_map['TAA'] = '*'
    codon_to_aa_map['TAG'] = '*'

    four_fold_deg(codon_to_aa_map, 'GC', 'A')
    
    codon_to_aa_map['TGC'] = 'C'
    codon_to_aa_map['TGT'] = 'C'
    codon_to_aa_map['TGT'] = 'C'
    codon_to_aa_map['TGY'] = 'C'
    
    codon_to_aa_map['GAC'] = 'D'
    codon_to_aa_map['GAT'] = 'D'
    codon_to_aa_map['GAY'] = 'D'
    
    codon_to_aa_map['GAA'] = 'E'
    codon_to_aa_map['GAG'] = 'E'
    codon_to_aa_map['GAR'] = 'E'
    
    codon_to_aa_map['TTC'] = 'F'
    codon_to_aa_map['TTT'] = 'F'
    codon_to_aa_map['TTY'] = 'F'
    
    four_fold_deg(codon_to_aa_map, 'GG', 'G')
    
    codon_to_aa_map['CAC'] = 'H'
    codon_to_aa_map['CAT'] = 'H'
    codon_to_aa_map['CAY'] = 'H'
    
    codon_to_aa_map['ATA'] = 'I'
    codon_to_aa_map['ATC'] = 'I'
    codon_to_aa_map['ATT'] = 'I'
    codon_to_aa_map['ATH'] = 'I'
    
    codon_to_aa_map['AAA'] = 'K'
    codon_to_aa_map['AAG'] = 'K'
    codon_to_aa_map['AAG'] = 'K'
    
    four_fold_deg(codon_to_aa_map, 'CT', 'L')
    codon_to_aa_map['TTA'] = 'L'
    codon_to_aa_map['TTG'] = 'L'
    codon_to_aa_map['TTR'] = 'L'
    codon_to_aa_map['YTR'] = 'L'

    codon_to_aa_map['ATG'] = 'M'
    
    codon_to_aa_map['AAC'] = 'N'
    codon_to_aa_map['AAT'] = 'N'
    codon_to_aa_map['AAY'] = 'N'

    four_fold_deg(codon_to_aa_map, 'CC', 'P')


    codon_to_aa_map['CAA'] = 'Q'
    codon_to_aa_map['CAG'] = 'Q'
    codon_to_aa_map['CAR'] = 'Q'
    
    
    four_fold_deg(codon_to_aa_map, 'CG', 'R')
    codon_to_aa_map['AGA'] = 'R'
    codon_to_aa_map['AGG'] = 'R'
    codon_to_aa_map['AGR'] = 'R'
    codon_to_aa_map['MGR'] = 'R'
    
    
    
    four_fold_deg(codon_to_aa_map, 'TC', 'S')
    codon_to_aa_map['AGC'] = 'S'
    codon_to_aa_map['AGT'] = 'S'
    codon_to_aa_map['WGY'] = 'S'

    four_fold_deg(codon_to_aa_map, 'AC', 'T')
    
    four_fold_deg(codon_to_aa_map, 'GT', 'V')
    
    codon_to_aa_map['TGG'] = 'W'
    
    codon_to_aa_map['TAC'] = 'Y'
    codon_to_aa_map['TAT'] = 'Y'
    codon_to_aa_map['TAY'] = 'Y'

    codon_to_aa_map['---'] = '-'
    codon_to_aa_map['...'] = '.'

    return codon_to_aa_map


#_________________________________________________________________________________________

#   Make C-string map

#_________________________________________________________________________________________
def create_char_array_map(codon_to_aa_map_python_dict):
    #sys.stderr.write("Create lookup table..\n")
    codon_to_aa_map_char_array = 'X' * 256*256*256
    #sys.stderr.write("prepare lookup table..\n")

    cdef:
        char * pcodon_to_aa_map_char_array = codon_to_aa_map_char_array
        char* kk
        char* vv
        int a, b, c
        str k, v

    for k,v in codon_to_aa_map_python_dict.iteritems():
        kk = k
        vv = v
        a = kk[0]
        b = kk[1]
        c = kk[2]
        pcodon_to_aa_map_char_array[a*256*256 + b * 256 + c] = vv[0]

    return codon_to_aa_map_char_array

    
    








codon_to_aa_map_python_dict = create_codon_to_aa_map()
codon_to_aa_map_char_array = create_char_array_map(codon_to_aa_map_python_dict)
codon_to_aa_map = codon_to_aa_map_python_dict








##########################################################################################

#   translate using python dict

##########################################################################################
def slow_translate (cdna):
    #sys.stderr.write("get peptide..")
    peptide = [codon_to_aa_map_python_dict.get(cdna[x:x+3], 'X')
                        for x in xrange(0, len(cdna), 3) ]
    #sys.stderr.write("join..")
    peptide_str =  "".join(peptide)
    #sys.stderr.write("done\n")
    return peptide_str


import cython


##########################################################################################

#   translate using char array

##########################################################################################
@cython.boundscheck(False)
def translate (str orig_cdna):

    cdef:
        char * pcodon_to_aa_map_char_array = codon_to_aa_map_char_array

    cdna_len = len(orig_cdna) - 2
    peptide = " " * ((cdna_len+2) // 3)

    #sys.stderr.write("loop..\n")

    cdef:
        char* cdna = orig_cdna
        int i
        int cnt_aa = (cdna_len+2) // 3
        char* ppeptide = peptide
        int x
        int a, b, c
    
    for i in range(cnt_aa):
        x = i * 3
        a = cdna[x]
        b = cdna[x+1]
        c = cdna[x+2]
        # faster here
        ppeptide[i] = pcodon_to_aa_map_char_array[a * 65536 + b * 256 + c]
    #sys.stderr.write("Done!\n")
    return peptide


    
 
 








##88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
#
## numpy
#
##88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
#    
##_________________________________________________________________________________________
#
##   Make numpy map
#
##_________________________________________________________________________________________
#cimport numpy as np
#import numpy as np
#ctypedef np.int8_t DTYPE_T
#def create_numpy_map(codon_to_aa_map_python_dict):
#
#    cdef np.ndarray[DTYPE_T, ndim=1] numpy_map = np.zeros([256*256*256], dtype =np.int8)
#    numpy_map.flat = ord('X')
#    cdef str k
#    cdef str v
#    cdef char* kk
#    cdef char* vv
#
#
#    for k,v in codon_to_aa_map_python_dict.iteritems():
#        kk = k
#        vv = v
#        numpy_map[kk[0] * 65536 + kk[1] * 256 + kk[2]] = vv[0]
#
#    return numpy_map
#
#
#codon_to_aa_map_numpy = create_numpy_map(codon_to_aa_map_python_dict)
#
#
##_________________________________________________________________________________________
#
##   translate using numpy
#
##_________________________________________________________________________________________
#import time
#import sys 
#
#
#
#
#@cython.boundscheck(False)
#def less_slow_translate (str orig_cdna):
#    cdef np.ndarray[DTYPE_T, ndim=1] numpy_map = codon_to_aa_map_numpy
#
#    cdef:
#        char * pcodon_to_aa_map_char_array = codon_to_aa_map_char_array
#
#    cdna_len = len(orig_cdna) - 2
#    peptide = " " * ((cdna_len+2) // 3)
#
#    cdef:
#        char* cdna = orig_cdna
#        unsigned int i
#        unsigned int cnt_aa = (cdna_len+2) // 3
#        char* ppeptide = peptide
#        unsigned int x
#        unsigned int a, b, c
#    for i in range(cnt_aa):
#        x = i * 3
#        a = cdna[x]
#        b = cdna[x+1]
#        c = cdna[x+2]
#        # slower here
#        ppeptide[i] = numpy_map[a * 65536 + b*256 + c]
#        #ppeptide[i] = pcodon_to_aa_map_char_array[a * 65536 + b * 256 + c]
#    return peptide
#
